package com.peng.leetcode.trie;

import java.util.*;

/**
 * TopKFrequent
 *
 * @author: lupeng6
 * @create: 2021/5/8 14:35
 */
public class TopKFrequent {

    public static void main(String[] args) {
        String[] strs = {"i", "love", "com/peng/leetcode", "i", "love", "coding"};
        List<String> list = new TopKFrequent().topKFrequent(strs, 2);
        System.out.println(list);
    }

    public List<String> topKFrequent(String[] words, int k) {
        Map<String, Integer> countMap = new HashMap<>();
        for (String word : words) {
            countMap.putIfAbsent(word, 0);
            countMap.compute(word, (key, v) -> v + 1);
        }

        PriorityQueue<String> queue = new PriorityQueue<String>((a, b) -> {
            return countMap.get(a).equals(countMap.get(b)) ? b.compareTo(a) : countMap.get(a).compareTo(countMap.get(b));
        });

        for (Map.Entry<String, Integer> entry : countMap.entrySet()) {
            queue.offer(entry.getKey());
            if (queue.size() > k) {
                queue.poll();
            }
        }

        List<String> ans = new ArrayList<>();

        while (!queue.isEmpty()) {
            ans.add(queue.poll());
        }
        Collections.reverse(ans);

        return ans;
    }

}
